module solid_rotation_test_mod

	use const_mod
	use mesh_mod
	use state_mod
	use tracer_mod
	use namelist_mod
	use sphere_geometry_mod
	use parallel_mod

	implicit none

	private
	public solid_rotation_test_init
	public solid_rotation_test_set_ic
	public solid_rotation_test_set_uv

	real(r8), parameter :: period = 12 * 86400 ! s
	real(r8), parameter :: h0     = 1000 ! m
	real(r8), parameter :: lon0   = 3 * pi / 2.0_r8
	real(r8), parameter :: lat0   = 0.0_r8
	real(r8), parameter :: alpha  = 90.0_r8 * rad
	real(r8), parameter :: lon1   = 3 * pi / 2.0_r8
	real(r8), parameter :: lat1   = 0.0_r8
	real(r8) u0

contains
	
	subroutine solid_rotation_test_init()

		u0 = pi2 * radius / period
	    
	end subroutine solid_rotation_test_init

	subroutine solid_rotation_test_set_ic(mesh, tracer)

		type(mesh_type), intent(in) :: mesh
		type(tracer_type), intent(inout) :: tracer
		integer i, j, itracer
    real(r8) lon, lat, r, R0, x1(3), x(3)

    if (ntracers /= 3) then
    	print*, 'solid_rotation_test case need 3 tracers!'
    	stop
    end if

    R0 = radius / 3.0_r8
    ! Background
    tracer%q(:,:,:,1) = 1.d0
		! cosine bell
		do j = mesh%full_lat_ibeg, mesh%full_lat_iend
			lat = mesh%full_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				r = calc_distance(lon0, lat0, lon, lat)
				if (r < R0) then
					tracer%q(i,j,1,2) = h0 / 2.0_r8 * (1 + cos(pi * r / R0))
				else
					tracer%q(i,j,1,2) = 0
				end if
			end do
		end do
		! Gaussian bell
		call cartesian_transform(lon1, lat1, x1(1), x1(2), x1(3)); x1 = x1 / radius
		do j = mesh%full_lat_ibeg, mesh%full_lat_iend
			lat = mesh%full_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				call cartesian_transform(lon, lat, x(1), x(2), x(3)); x = x / radius
				tracer%q(i,j,1,3) = exp(-5.0 * dot_product(x - x1, x - x1))
			end do
		end do

		do itracer = 1, size(tracer%q, 4)
			call fill_zonal_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
			call fill_merid_halo_cell(tracer%q(:,:,:,itracer), all_halo=.true.)
		end do

	end subroutine solid_rotation_test_set_ic

	subroutine solid_rotation_test_set_uv(mesh, state, elapsed_seconds)

		type(mesh_type), intent(in) :: mesh
		type(state_type), intent(inout) :: state
		real(8), intent(in) :: elapsed_seconds

		integer i, j, k
		real(r8) lon, lat

		associate (u => state%u, v => state%v, psi => state%psi)
		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	lat = mesh%half_lat(j)
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		lon = mesh%half_lon(i)
		! 		psi(i,j) = radius * u0 * (sin(lat) * cos(alpha) - cos(lon) * cos(lat) * sin(alpha))
		! 	end do
		! end do
		! call fill_zonal_halo_lat(psi, all_halo=.true.)

		! do j = mesh%full_lat_ibeg + 1, mesh%full_lat_iend - 1
		! 	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
		! 		u(i,j,1) = (psi(i,j) - psi(i,j-1)) / mesh%le_lon(j)
		! 	end do
		! end do
  !   call fill_zonal_halo_lon(u, all_halo=.true.)
		
		! do j = mesh%half_lat_ibeg, mesh%half_lat_iend
		! 	do i = mesh%full_lon_ibeg, mesh%full_lon_iend
		! 		v(i,j,1) = - (psi(i,j) - psi(i-1,j)) / mesh%le_lat(j)
		! 	end do
		! end do

	  do j = mesh%full_lat_ibeg, mesh%full_lat_iend
	  	lat = mesh%full_lat(j)
	  	do i = mesh%half_lon_ibeg, mesh%half_lon_iend
	  		lon = mesh%half_lon(i)
	  		u(i,j,1) = u0 * (cos(lat) * cos(alpha) + sin(lat) * cos(lon) * sin(alpha))
	  	end do
	  end do
		call fill_zonal_halo_lon(u, all_halo=.true.)

		do j = mesh%half_lat_ibeg, mesh%half_lat_iend
			lat = mesh%half_lat(j)
			do i = mesh%full_lon_ibeg, mesh%full_lon_iend
				lon = mesh%full_lon(i)
				v(i,j,1) = -u0 * sin(lon) * sin(alpha)
			end do
		end do
		end associate

	end subroutine solid_rotation_test_set_uv

end module solid_rotation_test_mod